import torch
import torch.utils.data as Data
from torchvision import transforms
from torchvision.datasets import FashionMNIST
from model import GoogLeNet, Inception
from torch import nn
import pandas as pd
import time


def test_data_process():
    test_data = FashionMNIST(root="./data",
                             train=False,
                             transform=transforms.Compose([transforms.Resize(size=224), transforms.ToTensor()]),
                             download=True)
    """
    names = ""
    for label, name in enumerate(test_data.classes):
        #print(f"标签 {label}: {name}")
        names += f"'{name}',"
    print("FashionMNIST数据集类别：", names)
    """
    test_dataloader = Data.DataLoader(dataset=test_data,
                                      batch_size=1,
                                      shuffle=True,
                                      num_workers=0)
    return test_dataloader


def test_model_process(model, test_dataloader):
    # 设定训练所用到的设备，有GPU用GPU没有GPU用CPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 将模型放入到训练设备中
    model = model.to(device)

    # 初始化参数
    test_corrects = 0
    test_num = 0
    # 只进行前向传播， 不计算梯度，从而节省内存，加快运行速度
    with torch.no_grad():
        for (test_data_x, test_data_y) in test_dataloader:
            test_data_x = test_data_x.to(device)
            test_data_y = test_data_y.to(device)
            model.eval()
            # 前向传播过程，输入为测试数据集，输出为每个样式的预测值
            output = model(test_data_x)
            # 查找每一行中最大值对应的下标
            pre_lab = torch.argmax(output, dim=1)
            # 预测正确，准确度test_corrects+1
            test_corrects += torch.sum(pre_lab == test_data_y.data)
            test_num += test_data_x.size(0)
    # 计算准确率
    test_acc = test_corrects.double().item() / test_num
    print("测试的准确率为：", test_acc)


if (__name__ == "__main__"):
    model = GoogLeNet(Inception)
    model.load_state_dict(torch.load("./model/best_model.pth"))
    test_dataloader = test_data_process()

    # test_model_process(model, test_dataloader)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    with torch.no_grad():
        # 初始化参数
        test_corrects = 0
        test_num = 0
        classes_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag',
                         'Ankle boot']
        for b_x, b_y in test_dataloader:
            b_x = b_x.to(device)
            b_y = b_y.to(device)
            # 设置为验证（评估模式）
            model.eval()
            output = model(b_x)
            pre_lab = torch.argmax(output, dim=1)
            result = pre_lab.item()
            label = b_y.item()
            print("预测值：", classes_names[result], "---------", "真实值：", classes_names[label])
            # 预测正确，准确度test_corrects+1
            test_corrects += torch.sum(pre_lab == b_y.data)
            test_num += b_x.size(0)
        # 计算准确率
        test_acc = test_corrects.double().item() / test_num
        print("测试的准确率为：", test_acc)
